package com.fan.chat.listener;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.fan.chat.model.JSONResult;
import com.plexpt.chatgpt.listener.AbstractStreamListener;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StringUtils;

import javax.websocket.Session;
import java.io.IOException;

@Slf4j
public class WebsocketStreamListener extends AbstractStreamListener {

    private final Session session;

    public WebsocketStreamListener(Session session) {
        this.session = session;
        super.setOnComplate(msg -> {
            log.info("request response msg: {}", msg);
            sendText(doneResult());
        });
    }

    @Override
    public void onMsg(String message) {
        sendText(message);
    }

    @Override
    public void onError(Throwable throwable, String response) {
        log.info("response error: {}", response);
        try {
            if (StringUtils.hasLength(response)) {
                JSONObject json = JSON.parseObject(response);
                response = "```json\n" + JSON.toJSONString(json, SerializerFeature.PrettyFormat) + "\n```";
            }
        } catch (Exception ignore) {
        }
        sendText(response);
        sendText(doneResult());
    }

    private static String doneResult() {
        return "[DONE]";
    }

    private void sendText(String message) {
        try {
            session.getBasicRemote().sendText(JSONResult.success(message).toString());
        } catch (IOException e) {
            log.info("Websocket send message is error", e);
        }
    }
}
